import gradio as gr
import os
import json
import time
import shutil
import uuid
from typing import Optional, Tuple, Dict, Any

from Core.comfy_client import ComfyUIClient
from Core.baidu_translate import baidu_translate


class LoRATestInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        self.workflow_path = os.path.join(
            os.path.dirname(__file__), "..", "Comfyui", "work_flow", "LoRA测试.json"
        )

        # 目录：loras、input、temp
        from config.config import get_lora_dir_path, get_input_path
        self.lora_dir = get_lora_dir_path()
        self.input_dir = get_input_path()
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)
        os.makedirs(self.input_dir, exist_ok=True)

        # 收集 LoRA 选项（包含子目录）
        self.lora_choices = self._collect_lora_choices()
        self.default_lora = "test\\ETShoeF.safetensors"
        if self.lora_choices and self.default_lora not in self.lora_choices:
            self.default_lora = self.lora_choices[0]

    def _collect_lora_choices(self):
        choices = []
        try:
            if os.path.exists(self.lora_dir):
                for root, _, files in os.walk(self.lora_dir):
                    for f in files:
                        if f.lower().endswith((".safetensors", ".ckpt")):
                            rel = os.path.relpath(os.path.join(root, f), self.lora_dir)
                            rel = rel.replace("/", "\\")
                            choices.append(rel)
        except Exception as e:
            print(f"收集LoRA失败: {e}")
        choices.sort()
        return choices

    def load_workflow(self) -> Optional[Dict[str, Any]]:
        try:
            with open(self.workflow_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"加载工作流失败: {e}")
            return None

    def _copy_image_to_input(self, src_path: str, prefix: str) -> Tuple[Optional[str], Optional[str]]:
        try:
            if not src_path or not os.path.exists(src_path):
                return None, "错误：未提供有效的图像路径"
            ext = os.path.splitext(src_path)[1].lower()
            if ext not in [".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"]:
                return None, "错误：仅支持PNG/JPG/JPEG/BMP/WEBP/TIFF"
            filename = f"{prefix}_{uuid.uuid4().hex}{ext}"
            dst = os.path.join(self.input_dir, filename)
            shutil.copy2(src_path, dst)
            return filename, None
        except Exception as e:
            return None, f"复制图片失败：{e}"

    def run_workflow(
        self,
        input_image_path: Optional[str],
        user_text: str,
        width: int,
        height: int,
        lora1_name: str,
        steps1: int,
        denoise1: float,
        lora2_name: str,
        steps2: int,
        denoise2: float,
        lora3_name: str,
        steps3: int,
        denoise3: float,
        lora4_name: str,
        steps4: int,
        denoise4: float,
        max_retries: int = 150,
        retry_interval: float = 2.0,
    ):
        try:
            workflow = self.load_workflow()
            if not workflow:
                return None, None, None, None, "无法加载工作流", None

            # 1) 文本翻译：节点67 TextInput_
            try:
                english = baidu_translate(user_text or "")
                translate_info = "翻译成功"
            except Exception as e:
                english = user_text or ""
                translate_info = f"翻译失败，改用原文本: {str(e)}"
            if "67" in workflow and "inputs" in workflow["67"]:
                workflow["67"]["inputs"]["text"] = english

            # 2) 节点62：加载图片（提供给 Joy Caption 侧链，不影响LoRA生成）
            if input_image_path:
                filename, err = self._copy_image_to_input(input_image_path, "lora_test")
                if err:
                    return None, None, None, None, f"生成失败：{err}", english
                if "62" in workflow and "inputs" in workflow["62"]:
                    workflow["62"]["inputs"]["image"] = filename

            # 3) 节点15：EmptyLatent 宽高
            if "15" in workflow and "inputs" in workflow["15"]:
                workflow["15"]["inputs"]["width"] = int(width)
                workflow["15"]["inputs"]["height"] = int(height)

            # 4) 四路 LoRA 与 K采样器设置
            # LoRA1 节点17 + K采样器10
            if "17" in workflow and "inputs" in workflow["17"]:
                workflow["17"]["inputs"]["lora_name"] = lora1_name
            if "10" in workflow and "inputs" in workflow["10"]:
                workflow["10"]["inputs"]["steps"] = int(steps1)
                workflow["10"]["inputs"]["denoise"] = float(denoise1)

            # LoRA2 节点27 + K采样器36
            if "27" in workflow and "inputs" in workflow["27"]:
                workflow["27"]["inputs"]["lora_name"] = lora2_name
            if "36" in workflow and "inputs" in workflow["36"]:
                workflow["36"]["inputs"]["steps"] = int(steps2)
                workflow["36"]["inputs"]["denoise"] = float(denoise2)

            # LoRA3 节点49 + K采样器47
            if "49" in workflow and "inputs" in workflow["49"]:
                workflow["49"]["inputs"]["lora_name"] = lora3_name
            if "47" in workflow and "inputs" in workflow["47"]:
                workflow["47"]["inputs"]["steps"] = int(steps3)
                workflow["47"]["inputs"]["denoise"] = float(denoise3)

            # LoRA4 节点61 + K采样器59
            if "61" in workflow and "inputs" in workflow["61"]:
                workflow["61"]["inputs"]["lora_name"] = lora4_name
            if "59" in workflow and "inputs" in workflow["59"]:
                workflow["59"]["inputs"]["steps"] = int(steps4)
                workflow["59"]["inputs"]["denoise"] = float(denoise4)

            # 5) 提交生成请求
            response = self.client.post_prompt(workflow)
            prompt_id = response.get("prompt_id")
            if not prompt_id:
                return None, None, None, None, "生成失败：未获得prompt_id", english

            # 6) 轮询输出，分别捕获四路 SaveImage 的产物
            result1 = None
            result2 = None
            result3 = None
            result4 = None
            status = "等待输出..."

            for _ in range(max_retries):
                history = self.client.get_history(prompt_id)
                try:
                    if prompt_id in history and "outputs" in history[prompt_id]:
                        outputs = history[prompt_id]["outputs"]

                        def _try_get_image_from_node(node_id: str) -> Optional[str]:
                            if node_id not in outputs:
                                return None
                            data = outputs[node_id]
                            if "images" in data and data["images"]:
                                image = data["images"][0]
                                try:
                                    img_bytes = self.client.get_image(image["filename"], image.get("subfolder"), image.get("type"))
                                    local = os.path.join(self.temp_dir, f"lora_test_{node_id}_{uuid.uuid4().hex}.png")
                                    with open(local, "wb") as f:
                                        f.write(img_bytes)
                                    return os.path.abspath(local)
                                except Exception:
                                    return None
                            return None

                        # SaveImage 节点：12, 35, 46, 58
                        result1 = result1 or _try_get_image_from_node("12")
                        result2 = result2 or _try_get_image_from_node("35")
                        result3 = result3 or _try_get_image_from_node("46")
                        result4 = result4 or _try_get_image_from_node("58")

                        if result1 or result2 or result3 or result4:
                            status = "生成完成\n" + translate_info
                            return result1, result2, result3, result4, status, english
                except Exception:
                    pass
                time.sleep(retry_interval)

            return None, None, None, None, "生成失败：超时或无输出", english
        except Exception as e:
            print(f"生成异常: {e}")
            return None, None, None, None, f"生成失败：{e}", None


def Tab_LoRA_test():
    interface = LoRATestInterface()

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### LoRA测试（四路对比生成）")

            # 节点62：加载图片
            load_image = gr.Image(label="加载图片(节点62)", type="filepath", interactive=True)

            # 节点67：文本输入（翻译到英文）
            text_input = gr.Textbox(label="文本(中文自动翻译)", lines=3)
            translated_text = gr.Textbox(label="英文翻译(用于工作流)", lines=3, interactive=False)

            # 节点15：EmptyLatent 宽高
            width_input = gr.Number(label="宽度(width)", value=1024, minimum=0, maximum=2048)
            height_input = gr.Number(label="高度(height)", value=1024, minimum=0, maximum=2048)

            # LoRA1：加载器与采样器
            lora1_dropdown = gr.Dropdown(
                label="加载LoRA1模型(节点17)",
                choices=interface.lora_choices,
                value=interface.default_lora,
                interactive=bool(interface.lora_choices),
            )
            steps1_input = gr.Number(label="LoRA1采样器(steps)", value=30, minimum=0, maximum=60)
            denoise1_slider = gr.Slider(label="LoRA1采样器(denoise)", minimum=0, maximum=1, value=1, step=0.01)

            # LoRA2
            lora2_dropdown = gr.Dropdown(
                label="加载LoRA2模型(节点27)",
                choices=interface.lora_choices,
                value=interface.default_lora,
                interactive=bool(interface.lora_choices),
            )
            steps2_input = gr.Number(label="LoRA2采样器(steps)", value=30, minimum=0, maximum=60)
            denoise2_slider = gr.Slider(label="LoRA2采样器(denoise)", minimum=0, maximum=1, value=1, step=0.01)

            # LoRA3
            lora3_dropdown = gr.Dropdown(
                label="加载LoRA3模型(节点49)",
                choices=interface.lora_choices,
                value=interface.default_lora,
                interactive=bool(interface.lora_choices),
            )
            steps3_input = gr.Number(label="LoRA3采样器(steps)", value=30, minimum=0, maximum=60)
            denoise3_slider = gr.Slider(label="LoRA3采样器(denoise)", minimum=0, maximum=1, value=1, step=0.01)

            # LoRA4
            lora4_dropdown = gr.Dropdown(
                label="加载LoRA4模型(节点61)",
                choices=interface.lora_choices,
                value=interface.default_lora,
                interactive=bool(interface.lora_choices),
            )
            steps4_input = gr.Number(label="LoRA4采样器(steps)", value=30, minimum=0, maximum=60)
            denoise4_slider = gr.Slider(label="LoRA4采样器(denoise)", minimum=0, maximum=1, value=1, step=0.01)

        with gr.Column(scale=2):
            generate_btn = gr.Button("生成", variant="primary")
            status_text = gr.Textbox(label="状态", interactive=False)

            with gr.Row():
                result_image1 = gr.Image(label="LoRA1 生成预览")
                result_image2 = gr.Image(label="LoRA2 生成预览")
            with gr.Row():
                result_image3 = gr.Image(label="LoRA3 生成预览")
                result_image4 = gr.Image(label="LoRA4 生成预览")

    def translate_and_run(
        img_path,
        text,
        width,
        height,
        lora1,
        steps1,
        denoise1,
        lora2,
        steps2,
        denoise2,
        lora3,
        steps3,
        denoise3,
        lora4,
        steps4,
        denoise4,
    ):
        r1, r2, r3, r4, status, english = interface.run_workflow(
            img_path,
            text,
            int(width or 1024),
            int(height or 1024),
            lora1,
            int(steps1 or 30),
            float(denoise1 or 1),
            lora2,
            int(steps2 or 30),
            float(denoise2 or 1),
            lora3,
            int(steps3 or 30),
            float(denoise3 or 1),
            lora4,
            int(steps4 or 30),
            float(denoise4 or 1),
        )
        return r1, r2, r3, r4, status, english

    generate_btn.click(
        fn=translate_and_run,
        inputs=[
            load_image,
            text_input,
            width_input,
            height_input,
            lora1_dropdown,
            steps1_input,
            denoise1_slider,
            lora2_dropdown,
            steps2_input,
            denoise2_slider,
            lora3_dropdown,
            steps3_input,
            denoise3_slider,
            lora4_dropdown,
            steps4_input,
            denoise4_slider,
        ],
        outputs=[
            result_image1,
            result_image2,
            result_image3,
            result_image4,
            status_text,
            translated_text,
        ],
    )